package org.ehe.ai.service;

import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import lombok.extern.slf4j.Slf4j;
import org.ehe.ai.domain.ChatMessage;
import org.ehe.ai.domain.ChatSession;
import org.ehe.ai.domain.dto.ChatRequest;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.List;
import java.util.Objects;
import java.util.concurrent.*;
import java.util.regex.Pattern;

/**
 * @author :zhangnn
 * @className :StreamChatService
 * @description: TODO
 * @date 2025-08-23 09:34:43
 */
@Service
@Slf4j
public class StreamChatService {

    @Autowired
    private ChatService chatService;

    //存储每个会话的心跳任务，用于取消
    private final ConcurrentHashMap<String, ScheduledFuture<?>> heartbeatTasks = new ConcurrentHashMap<>();

    //跟踪emitter是否已完成的状态
    private final ConcurrentHashMap<String, Boolean> emitterCompleted = new ConcurrentHashMap<>();

    // 存储活跃的SSE连接
    private final ConcurrentHashMap<String, SseEmitter> activeEmitters = new ConcurrentHashMap<>();

    // 线程池用于异步处理
    private final ScheduledExecutorService executorService = Executors.newScheduledThreadPool(20);

    // 中文分词规则
    private static final Pattern CHINESE_SENTENCE_PATTERN = Pattern.compile("[。！？；：]");
    private static final Pattern CHINESE_PHRASE_PATTERN = Pattern.compile("[，、]");

    /**
     * 创建SSE连接
     */
    public SseEmitter createSseConnection(String sessionId) {
        // 强制清理可能残留的旧连接状态
        closeSseConnection(sessionId); // 调用现有方法关闭旧连接

        // 设置超时时间为30分钟
        SseEmitter emitter = new SseEmitter(30 * 60 * 1000L);
        emitterCompleted.put(sessionId, false); // 初始化新连接状态


        // 连接建立时发送欢迎消息
        try {
            emitter.send(SseEmitter.event()
                .name("connected")
                .data("连接建立成功，会话ID: " + sessionId)
                .id(String.valueOf(System.currentTimeMillis())));

            log.info("[AI]SSE连接建立成功: {}", sessionId);
        } catch (IOException e) {
            emitterCompleted.put(sessionId, true);
            log.error("发送连接成功消息失败: {}", sessionId, e);
            return null;
        }


        activeEmitters.put(sessionId, emitter);
        // 存储连接
//        SseEmitter existingEmitter = activeEmitters.put(sessionId, emitter);
//        if (existingEmitter != null) {
//            try {
//                existingEmitter.complete();
//                log.info("关闭已存在的SSE连接: {}", sessionId);
//            } catch (Exception e) {
//                log.warn("关闭已存在连接失败", e);
//            }
//        }

        // 设置连接生命周期处理
        setupEmitterLifecycle(sessionId, emitter);

        // 启动心跳
        scheduleHeartbeat(sessionId, emitter);

        return emitter;
    }

    /**
     * 设置SSE连接生命周期处理
     */
    private void setupEmitterLifecycle(String sessionId, SseEmitter emitter) {
        emitter.onCompletion(() -> {
            log.info("SSE连接完成: {}", sessionId);
            activeEmitters.remove(sessionId);
            emitterCompleted.put(sessionId, true); // 标记为已完成
            // 新增：取消心跳任务
            cancelHeartbeat(sessionId);
        });

        emitter.onTimeout(() -> {
            log.info("SSE连接超时: {}", sessionId);
            activeEmitters.remove(sessionId);
            emitterCompleted.put(sessionId, true); // 标记为已完成
            try {
                emitter.complete();
            } catch (Exception e) {
                log.warn("完成超时连接失败", e);
            }
            // 新增：取消心跳任务
            cancelHeartbeat(sessionId);
        });

        emitter.onError(throwable -> {
            log.error("SSE连接错误: {}", sessionId, throwable);
            activeEmitters.remove(sessionId);
            emitterCompleted.put(sessionId, true); // 标记为已完成
            try {
                emitter.complete();
            } catch (Exception e) {
                log.warn("完成错误连接失败", e);
            }
            // 新增：取消心跳任务
            cancelHeartbeat(sessionId);
        });
    }

    /**
     * 流式发送消息
     */
    public void streamMessage(ChatRequest request) {
        String sessionId = request.getSessionId();
        SseEmitter emitter = activeEmitters.get(sessionId);

        if (emitter == null) {
            ChatSession session = chatService.selectSession(sessionId);
            if(Objects.nonNull(session)){
                createSseConnection(sessionId);
                emitter = activeEmitters.get(sessionId);
                log.warn("重新连接并获取emitter===》: {}", sessionId);
            } else {
                log.warn("未找到会话的SSE连接进行再次连接: {}", sessionId);
                return;
            }
        }

        // 校验emitter是否有效（防止创建失败的情况）
        if (emitter == null) {
            log.error("创建连接后仍未获取到emitter: {}", sessionId);
            return;
        }

        // 异步处理AI请求
        SseEmitter finalEmitter = emitter;
        CompletableFuture.runAsync(() -> {
            processStreamMessage(request, finalEmitter);
        }, executorService);
    }

    /**
     * 处理流式消息
     */
    private void processStreamMessage(ChatRequest request, SseEmitter emitter) {
        log.info("进入到六十消息处理方法====processStreamMessage");
        String sessionId = request.getSessionId();
        long startTime = System.currentTimeMillis();

        try {
            // 验证或创建会话
            sessionId = chatService.validateOrCreateSession(sessionId, request.getUserId());
            request.setSessionId(sessionId);

            // 发送开始处理的消息（首次校验）
            SseEmitter currentEmitter = activeEmitters.get(sessionId);
            if (!isEmitterActive(sessionId, currentEmitter)) {
                log.error("初始连接已失效: {}", sessionId);
                return;
            }

            // 发送开始处理的消息
            sendEvent(emitter, "start", "正在思考中...", sessionId);

            // 保存用户消息
            chatService.saveMessage(sessionId, request.getMessage(), 1);
            //更新session标题
            int chats = chatService.countChatMessage(sessionId);
            log.info("获取chats访问数量={}",chats);
            if(chats < 2){
                chatService.updateSessionTitle(sessionId,request.getMessage());
            }
            // 调用AI API
            String aiResponse = chatService.callAiApi(request.getMessage());

            // 流式发送AI响应
            //streamAiResponse(emitter, aiResponse, sessionId);
            streamAiResponseWithFreshEmitter(sessionId, aiResponse);

            // 保存AI回复
            chatService.saveMessage(sessionId, aiResponse, 2);

            // 发送完成消息
            long processingTime = System.currentTimeMillis() - startTime;
            sendEvent(emitter, "complete",
                String.format("回答完成，耗时: %dms", processingTime), sessionId);

            log.info("流式消息处理完成，会话: {}, 耗时: {}ms", sessionId, processingTime);

        } catch (Exception e) {
            log.error("流式处理消息失败，会话: {}", sessionId, e);
            SseEmitter currentEmitter = activeEmitters.get(sessionId);
            sendEvent(currentEmitter, "error", "处理失败: " + e.getMessage(), sessionId);
        }
    }

    // 新增：带新鲜emitter的流式发送方法
    private void streamAiResponseWithFreshEmitter(String sessionId, String fullResponse) throws InterruptedException {
        if (fullResponse == null || fullResponse.isEmpty()) {
            SseEmitter emitter = activeEmitters.get(sessionId);
            sendEvent(emitter, "message", "暂无回复", sessionId);
            return;
        }

        String[] chunks = intelligentChunking(fullResponse);
        for (int i = 0; i < chunks.length; i++) {
            String chunk = chunks[i];
            if (chunk.trim().isEmpty()) continue;

            // 每次发送前都获取最新的emitter
            SseEmitter currentEmitter = activeEmitters.get(sessionId);
            if (!isEmitterActive(sessionId, currentEmitter)) {
                log.warn("连接已失效，停止发送: {}", sessionId);
                return;
            }

            sendEvent(currentEmitter, "message", chunk, sessionId);

            // 动态延迟
            int delay = calculateDelay(chunk, i, chunks.length);
            Thread.sleep(delay);
        }
    }
    /**
     * 智能流式发送AI响应
     */
    private void streamAiResponse(SseEmitter emitter, String fullResponse, String sessionId) {
        try {
            if (fullResponse == null || fullResponse.isEmpty()) {
                sendEvent(emitter, "message", "暂无回复", sessionId);
                return;
            }

            // 使用更智能的分块策略
            String[] chunks = intelligentChunking(fullResponse);

            for (int i = 0; i < chunks.length; i++) {
                String chunk = chunks[i];
                if (chunk.trim().isEmpty()) continue;

                // 新增：发送前检查连接是否有效
                if (!isEmitterActive(sessionId, emitter)) {
                    log.warn("连接已失效，停止发送: {}", sessionId);
                    return;
                }

                sendEvent(emitter, "message", chunk, sessionId);

                // 根据内容长度动态调整延迟
                int delay = calculateDelay(chunk, i, chunks.length);
                Thread.sleep(delay);
            }

        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            log.warn("流式发送被中断: {}", sessionId);
        } catch (Exception e) {
            log.error("流式发送响应失败: {}", sessionId, e);
            sendEvent(emitter, "error", "发送失败", sessionId);
        }
    }

    /**
     * 智能分块策略
     */
    private String[] intelligentChunking(String text) {
        // 首先按句子分割
        String[] sentences = CHINESE_SENTENCE_PATTERN.split(text);

        if (sentences.length > 1) {
            // 如果有多个句子，按句子分块
            String[] chunks = new String[sentences.length];
            for (int i = 0; i < sentences.length; i++) {
                chunks[i] = sentences[i] + (i < sentences.length - 1 ? getDelimiter(text, sentences[i]) : "");
            }
            return chunks;
        }

        // 单句子的情况，按短语分割
        String[] phrases = CHINESE_PHRASE_PATTERN.split(text);
        if (phrases.length > 1) {
            String[] chunks = new String[phrases.length];
            for (int i = 0; i < phrases.length; i++) {
                chunks[i] = phrases[i] + (i < phrases.length - 1 ? getDelimiter(text, phrases[i]) : "");
            }
            return chunks;
        }

        // 最后按字符数分割
        return chunkByCharacters(text, 15);
    }

    /**
     * 获取分隔符
     */
    private String getDelimiter(String originalText, String chunk) {
        int index = originalText.indexOf(chunk) + chunk.length();
        if (index < originalText.length()) {
            char nextChar = originalText.charAt(index);
            if ("。！？；：，、".indexOf(nextChar) != -1) {
                return String.valueOf(nextChar);
            }
        }
        return "";
    }

    /**
     * 按字符数分块
     */
    private String[] chunkByCharacters(String text, int chunkSize) {
        int length = text.length();
        int numChunks = (length + chunkSize - 1) / chunkSize;
        String[] chunks = new String[numChunks];

        for (int i = 0; i < numChunks; i++) {
            int start = i * chunkSize;
            int end = Math.min(start + chunkSize, length);
            chunks[i] = text.substring(start, end);
        }

        return chunks;
    }

    /**
     * 计算延迟时间
     */
    private int calculateDelay(String chunk, int index, int totalChunks) {
        // 基础延迟
        int baseDelay = 100;

        // 根据内容长度调整
        int lengthFactor = Math.min(chunk.length() * 8, 200);

        // 标点符号有更长的停顿
        if (chunk.matches(".*[。！？]$")) {
            lengthFactor += 300;
        } else if (chunk.matches(".*[，；：]$")) {
            lengthFactor += 150;
        }

        // 最后一块的延迟稍短
        if (index == totalChunks - 1) {
            lengthFactor = Math.min(lengthFactor, 100);
        }

        return baseDelay + lengthFactor + (int)(Math.random() * 50);
    }

    /**
     * 发送SSE事件
     */
    private void sendEvent(SseEmitter emitter, String eventName, String data, String sessionId) {
        if (!isEmitterActive(sessionId, emitter)) {
            log.debug("连接已失效，跳过发送事件: {}", sessionId);
            return;
        }

        try {
            emitter.send(SseEmitter.event()
                .name(eventName)
                .data(data)
                .id(String.valueOf(System.currentTimeMillis())));
        } catch (IOException e) {
            log.error("发送SSE事件失败: {} -> {}", sessionId, eventName, e);
            // 连接可能已断开，移除它
            activeEmitters.remove(sessionId);
            cancelHeartbeat(sessionId);
        }
    }

    /**
     * 发送心跳包
     */
    private void scheduleHeartbeat(String sessionId, SseEmitter emitter) {
        // 保存心跳任务到map
        ScheduledFuture<?> heartbeatTask = executorService.scheduleAtFixedRate(() -> {
            // 心跳发送前检查连接
            if (isEmitterActive(sessionId, emitter)) {
                try {
                    emitter.send(SseEmitter.event()
                        .name("heartbeat")
                        .data("ping")
                        .id(String.valueOf(System.currentTimeMillis())));
                } catch (IOException e) {
                    log.warn("心跳包发送失败，移除连接: {}", sessionId);
                    activeEmitters.remove(sessionId);
                    emitterCompleted.put(sessionId, true);
                    cancelHeartbeat(sessionId);
                }
            } else {
                // 连接已失效，主动取消任务
                cancelHeartbeat(sessionId);
            }
        }, 15, 15, TimeUnit.SECONDS);

        heartbeatTasks.put(sessionId, heartbeatTask);
    }

    /**
     * 关闭SSE连接
     */
    public void closeSseConnection(String sessionId) {
        SseEmitter emitter = activeEmitters.remove(sessionId);
        if (emitter != null) {
            try {
                emitter.send(SseEmitter.event()
                    .name("close")
                    .data("连接即将关闭"));
                emitter.complete();
                log.info("SSE连接已关闭: {}", sessionId);
            } catch (IOException e) {
                log.error("关闭SSE连接失败: {}", sessionId, e);
            } finally {
                // 确保心跳任务被取消
                cancelHeartbeat(sessionId);
            }
        }
    }

    /**
     * 获取活跃连接数
     */
    public int getActiveConnectionCount() {
        return activeEmitters.size();
    }

    /**
     * 清理所有连接（用于应用关闭时）
     */
    public void cleanupAllConnections() {
        log.info("清理所有SSE连接，当前连接数: {}", activeEmitters.size());

        activeEmitters.forEach((sessionId, emitter) -> {
            try {
                emitter.send(SseEmitter.event()
                    .name("shutdown")
                    .data("服务器即将关闭"));
                emitter.complete();
            } catch (Exception e) {
                log.warn("清理连接失败: {}", sessionId, e);
            } finally {
                cancelHeartbeat(sessionId);
            }
        });

        activeEmitters.clear();
        executorService.shutdown();
    }

    // 新增：辅助方法 - 校验连接是否活跃
    private boolean isEmitterActive(String sessionId, SseEmitter emitter) {
        if (emitter == null) {
            return false;
        }
        // 双重校验：内存中的emitter是否与当前存储的一致，且未被标记为完成
        SseEmitter currentStoredEmitter = activeEmitters.get(sessionId);
        Boolean completed = emitterCompleted.getOrDefault(sessionId, true); // 默认已完成
        return currentStoredEmitter != null
            && currentStoredEmitter == emitter
            && !completed;
    }

    // 新增：辅助方法 - 取消心跳任务
    private void cancelHeartbeat(String sessionId) {
        ScheduledFuture<?> task = heartbeatTasks.remove(sessionId);
        if (task != null && !task.isCancelled()) {
            task.cancel(false);
            log.debug("已取消会话[{}]的心跳任务", sessionId);
        }
    }
}
